/**
 * @file myallocator.h
 * @author like
 * @brief 移植SGI STL二级空间配置器内存池源码
 * @version 0.1
 * @date 2025-05-25
 * 
 * @copyright Copyright (c) 2025
 */

#pragma once

#include <cstddef>
#include <cstring>
#include <mutex>
#include <iostream>

// 封装了malloc和free操作，可以设置OOM释放内存的回调函数
template <int __inst>
class __malloc_alloc_template
{
private:
    static void* _S_oom_malloc(size_t);
    static void* _S_oom_realloc(void*, size_t);

    static void (* __malloc_alloc_oom_handler)();
public:
    static void* allocate(size_t __n)
    {
        void* __result = malloc(__n);
        if (0 == __result) __result = _S_oom_malloc(__n);
        return __result;
    }

    static void deallocate(void* __p, size_t /* __n */)
    {
        free(__p);
    }

    static void* reallocate(void* __p, size_t /* old_sz */, size_t __new_sz)
    {
        void* __result = realloc(__p, __new_sz);
        if (0 == __result) __result = _S_oom_realloc(__p, __new_sz);
        return __result;
    }

    static void (* __set_malloc_handler(void (*__f)()))()
    {
        void (* __old)() = __malloc_alloc_oom_handler;
        __malloc_alloc_oom_handler = __f;
        return(__old);
    }
};

template <int __inst>
void (* __malloc_alloc_template<__inst>::__malloc_alloc_oom_handler)() = nullptr;

template <int __inst>
void*
__malloc_alloc_template<__inst>::_S_oom_malloc(size_t __n)
{
    void (* __my_malloc_handler)();
    void* __result;

    for (;;) {
        __my_malloc_handler = __malloc_alloc_oom_handler;
        if (0 == __my_malloc_handler) { throw std::bad_alloc(); }
        (*__my_malloc_handler)();
        __result = malloc(__n);
        if (__result) return(__result);
    }
}

template <int __inst>
void* __malloc_alloc_template<__inst>::_S_oom_realloc(void* __p, size_t __n)
{
    void (* __my_malloc_handler)();
    void* __result;

    for (;;)
    {
        __my_malloc_handler = __malloc_alloc_oom_handler;
        if (0 == __my_malloc_handler) { throw std::bad_alloc(); }
        (*__my_malloc_handler)();
        __result = realloc(__p, __n);
        if (__result) return(__result);
    }
}

typedef __malloc_alloc_template<0> malloc_alloc;


/**
 * @brief 移植SGI STL二级空间配置器内存池
 * 
 * @tparam T 模板类型
 */
template<typename T>
class myallocator
{
public:
    using value_type = T;   // STL容器使用需要

    /**
     * @brief 向内存池申请内存
     * 
     * @param n 申请内存大小，单位：字节
     * @return void* 申请到内存的地址
     */
    T* allocate(size_t n)
    {
        n = n * sizeof(T);  // STL容器使用需要，STL容器中给空间配置器传入的是开辟对象的个数，而不是字节数
        //std::cout << __func__ << ": n = " <<  n << std::endl;
        void* ret = nullptr;

        if (n > (size_t) MAX_BYTES)
        {
            ret = malloc_alloc::allocate(n);
        }
        else
        {
            Obj* volatile* my_free_list
                = S_free_list_ + S_freelist_index(n);
            std::lock_guard<std::mutex> lock(mtx_);
            Obj* result = *my_free_list;
            if (result == nullptr)
            {
                ret = S_refill(S_round_up(n));
            }
            else
            {
                *my_free_list = result -> M_free_list_link_;
                ret = result;
            }
        }

        return (T*)ret;
    }


    /**
     * @brief 将内存归还到内存池中
     * 
     * @param p 待内存首地址
     * @param n 归还内存的大小
     */
    void deallocate(void* p, size_t n)
    {
        n = n * sizeof(T);  // STL容器使用需要，STL容器中给空间配置器传入的是释放对象的个数，而不是字节数
        //std::cout << __func__ << ": n = " <<  n << std::endl;
        if (n > (size_t) MAX_BYTES)
        {
            malloc_alloc::deallocate(p, n);
        }
        else
        {
            Obj* volatile*  my_free_list
                = S_free_list_ + S_freelist_index(n);
            Obj* q = (Obj*)p;

            std::lock_guard<std::mutex> lock(mtx_);
            q -> M_free_list_link_ = *my_free_list;
            *my_free_list = q;
            // lock is released here
        }
    }

    
    /**
     * @brief 内存扩容或者缩容
     * 
     * @param p 待扩容或缩容内存地址
     * @param old_sz 原内存大小，单位：字节
     * @param new_sz 新内存大小，单位：字节
     * @return void* 扩容或缩容后内存地址
     */
    void* reallocate(void* p, size_t old_sz, size_t new_sz)
    {
        //std::cout << __func__ << ": old_sz = " <<  old_sz
        //          << ", new_sz = " << new_sz << std::endl;
        void* result;
        size_t copy_sz;

        if (old_sz > (size_t)MAX_BYTES && new_sz > (size_t)MAX_BYTES)
        {
            return (realloc(p, new_sz));
        }

        if (S_round_up(old_sz) == S_round_up(new_sz))
        {
            return (p);
        }

        result = allocate(new_sz);
        copy_sz = new_sz > old_sz ? old_sz : new_sz;
        memcpy(result, p, copy_sz);
        deallocate(p, old_sz);
        return (result);
    }

    /**
     * @brief 对象构造
     * 
     * @tparam Args 可变参类型
     * @param p 构造对象的地址
     * @param args 构造对象的参数(可变参数)
     */
    template<typename... Args>
    void construct(T* p, Args&&... args)
    {
        new (p) T(std::forward<Args>(args)...);
    }

    ///**
    // * @brief 对象构造
    // * 
    // * @param p 构造对象的地址
    // * @param val 对象构造需要传入的参数
    // */
    //void construct(T* p, const T& val)
    //{
    //    new (p) T(val);
    //}

    /**
     * @brief 对象析构
     * 
     * @param p 待析构对象的地址
     */
    void destroy(T* p)
    {
        p->~T();
    }

private:
    enum {ALIGN = 8};              ///< 自由链表从8字节开始，以8字节对齐方式，一直扩充到128字节
    enum {MAX_BYTES = 128};        ///< 内存池最大的chunk块
    enum {NFREELISTS = 16};        ///< 自由链表的个数 MAX_BYTES/ALIGN

    /// @brief 每一个chunk块的头信息
    union Obj
    {
      union Obj* M_free_list_link_;     ///< 存储下一个chunk块的地址
      char M_client_data_[1];           ///< The client sees this.
    };

    static Obj* volatile S_free_list_[NFREELISTS];  ///< 存储自由链表数组的起始地址

    static char* S_start_free_;     ///< 表示在内存池中可以使用的内存的起始地址
    static char* S_end_free_;       ///< 表示在内存池中可以使用的内存的结束地址
    static size_t S_heap_size_;     ///< 表示内存池已经开辟的堆内存大小，单位：字节

    static std::mutex mtx_;         ///< 互斥锁，保证线程安全

    /**
     * @brief 将bytes字节上调到最临近ALIGN的倍数
     * 
     * @param bytes 待调整的字节数
     * @return size_t 调整后的字节数
     */
    static size_t S_round_up(size_t bytes)
    {
        return (((bytes) + (size_t)ALIGN - 1) & ~((size_t)ALIGN - 1));
    }

    /**
     * @brief 获取bytes字节对应在自由链表数组中的下标
     * 
     * @param bytes 待获取数组下标的字节数
     * @return size_t 传入字节数对应的在自由链表数组中的下标
     */
    static size_t S_freelist_index(size_t bytes)
    {
        return (((bytes) + (size_t)ALIGN - 1) / (size_t)ALIGN - 1);
    }

    /**
     * @brief 负责把分配好的chunk块进行连接，添加到自由链表当中
     * 
     * @param n 向内存池申请的内存大小
     * @return void* 申请的n字节大小的内存地址
     */
    static void* S_refill(size_t n)
    {
        int nobjs = 20;
        char* chunk = S_chunk_alloc(n, nobjs);
        Obj* volatile* my_free_list;
        Obj* result;
        Obj* current_obj;
        Obj* next_obj;
        int i;

        if (1 == nobjs) return(chunk);
        my_free_list = S_free_list_ + S_freelist_index(n);

        /* Build free list in chunk */
        result = (Obj*)chunk;
        *my_free_list = next_obj = (Obj*)(chunk + n);
        for (i = 1; ; i++)
        {
            current_obj = next_obj;
            next_obj = (Obj*)((char*)next_obj + n);
            if (nobjs - 1 == i)
            {
                current_obj -> M_free_list_link_ = 0;
                break;
            }
            else
            {
                current_obj -> M_free_list_link_ = next_obj;
            }
        }
        return(result);
    }

    /**
     * @brief 分配相应内存字节大小的chunk块
     * 
     * @param size 待分配的内存大小，单位：字节
     * @param nobjs 传入传出参数，分配chunk块的数量
     * @return char* 给用户使用的chunk块的内存地址
     */
    static char* S_chunk_alloc(size_t size, int& nobjs)
    {
        char* result;
        size_t total_bytes = size * nobjs;
        size_t bytes_left = S_end_free_ - S_start_free_;

        if (bytes_left >= total_bytes)
        {
            result = S_start_free_;
            S_start_free_ += total_bytes;
            return(result);
        }
        else if (bytes_left >= size)
        {
            nobjs = (int)(bytes_left/size);
            total_bytes = size * nobjs;
            result = S_start_free_;
            S_start_free_ += total_bytes;
            return(result);
        }
        else
        {
            size_t bytes_to_get = 
    	        2 * total_bytes + S_round_up(S_heap_size_ >> 4);
            // Try to make use of the left-over piece.
            if (bytes_left > 0)
            {
                Obj* volatile* my_free_list =
                    S_free_list_ + S_freelist_index(bytes_left);

                ((Obj*)S_start_free_) -> M_free_list_link_ = *my_free_list;
                *my_free_list = (Obj*)S_start_free_;
            }
            S_start_free_ = (char*)malloc(bytes_to_get);
            if (nullptr == S_start_free_)
            {
                size_t i;
                Obj* volatile* my_free_list;
    	        Obj* p;
                for (i = size;
                     i <= (size_t) MAX_BYTES;
                     i += (size_t) ALIGN)
                {
                    my_free_list = S_free_list_ + S_freelist_index(i);
                    p = *my_free_list;
                    if (0 != p)
                    {
                        *my_free_list = p -> M_free_list_link_;
                        S_start_free_ = (char*)p;
                        S_end_free_ = S_start_free_ + i;
                        return(S_chunk_alloc(size, nobjs));
                    }
                }
    	        S_end_free_ = 0;	// In case of exception.
                S_start_free_ = (char*)malloc_alloc::allocate(bytes_to_get);
            }
            S_heap_size_ += bytes_to_get;
            S_end_free_ = S_start_free_ + bytes_to_get;
            return(S_chunk_alloc(size, nobjs));
        }
    }

};

template<typename T>
typename myallocator<T>::Obj* volatile myallocator<T>::S_free_list_[NFREELISTS] =
{
    nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr,
    nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr 
};

template<typename T>
char* myallocator<T>::S_start_free_ = nullptr;

template<typename T>
char* myallocator<T>::S_end_free_ = nullptr;

template<typename T>
size_t myallocator<T>::S_heap_size_ = 0;

template<typename T>
std::mutex myallocator<T>::mtx_;
